import Data.List

----------------------------------------------------
----------------------------------------------------
-- Values for the DTMC
myMatrix :: [[Float]]
myMatrix = 	[[0.0, 1.0,  0.0,  0.0], 
			 [0.0, 0.01, 0.01, 0.98],
			 [1.0, 0.0,  0.0,  0.0],
			 [0.0, 0.0,  0.0,  1.0]]

estados :: [Char]
estados = "0123"

edoInicial :: Char
edoInicial = '0'

etiquetas :: [(Char, [String])]
etiquetas = [('0',[]), ('1',["try"]), ('2',["fail"]), ('3',["succ"])]
----------------------------------------------------
----------------------------------------------------


----------------------------------------------------
-- Data type for the state formulas
data Phi  = Verdad | Falso | Atom String | No Phi | Y Phi Phi | O Phi Phi | Imp Phi Phi | Prob String Float Psi
instance Show Phi where
	show = phiS

phiS :: Phi -> String
phiS Verdad = "T"
phiS Falso = "F"
phiS (Atom s) = s
phiS (No p) = "~" ++ phiS p
phiS (Y p1 p2) = (phiS p1) ++ ['^'] ++ (phiS p2)
phiS (O p1 p2) = (phiS p1) ++ ['v'] ++ (phiS p2)
phiS (Imp p1 p2) = (phiS p1) ++ "=>" ++ (phiS p2)
phiS (Prob com x p) = "P[" ++ com ++ (show x) ++ "](" ++ (show p) ++ ")"
----------------------------------------------------

----------------------------------------------------
-- Data type for the path formulas
data Psi = Next Phi | UntilB Int Phi Phi
instance Show Psi where
	show = psiS

psiS (Next p) = "X(" ++ show p ++ ")"
psiS (UntilB n p1 p2) = show p1 ++ "U[" ++ show n ++ "]" ++ show p2
----------------------------------------------------

----------------------------------------------------
-- Data type that represets the PCTL formulas 
data Pctl = Phi | Psi
----------------------------------------------------

----------------------------------------------------
-- Data type for DTMCs
data DTMC = DTMC {
						states	:: [Char],
						sInit	:: Char,
						pTrans	:: [[Float]],
						labelF	:: [(Char,[String])]
					} deriving (Show)

-- An instance of a DTMC named mChain
mChain = DTMC {states = estados, sInit = edoInicial, pTrans = myMatrix, labelF = etiquetas}
----------------------------------------------------

----------------------------------------------------
-- Te satisfiability algorithm
sat :: DTMC -> Phi -> String
sat m Verdad = states m
sat m Falso = []
sat m (Atom s) = [edo | edo<-(states m), lab<-(labelF m), edo==fst lab, elem s (snd lab)]
sat m (No p) = [edo | edo<-(states m), not (elem edo (sat m p))]
sat m (Y p1 p2) = intersect (sat m p1) (sat m p2)
sat m (O p1 p2) = union (sat m p1) (sat m p2)
sat m (Imp p1 p2) = union [edo | edo<-(states m), not(elem edo (sat m p1))] (sat m p2)
sat m (Prob com x p) 	|	com == "<"	= [fst pair | pair<-zip (states m) (ps m p), snd pair<x]
						|	com == "<="	= [fst pair | pair<-zip (states m) (ps m p), snd pair<=x]
						|	com == ">"	= [fst pair | pair<-zip (states m) (ps m p), snd pair>x]
						| 	com == ">="	= [fst pair | pair<-zip (states m) (ps m p), snd pair>=x]
						| 	otherwise	= error "Comparador no válido"
----------------------------------------------------

----------------------------------------------------
-- The ps function computes the probability vector for the states in the DTMC that satisfy 
-- the path formula 'p'
ps :: DTMC -> Psi -> [Float]
ps m (Next p) = next (pTrans m) (vecNext (states m) (sat m p))
ps m (UntilB 0 _ p2) = [if elem edo (sYes p2) then 1.0 else 0.0 | edo<-(states m)]
ps m (UntilB (n+1) p1 p2) = mult pPrime (ps m (UntilB n p1 p2))
							where 
								mult xss ys = [sum (zipWith (*) xs ys) | xs<-xss] 
								pPrime = [takeVal (fst myPair) (snd myPair) | myPair<-zip (states m) (pTrans m)]
								sY = sYes p2
								sN = sNo p1 p2
								sQ = [edo | edo<-(states m), not(elem edo (sY++sN))]
								takeVal st list  = if (elem st sQ) 
														then list
														else if (elem st sY)
																then [if e==st then 1.0 else 0.0 | e<-(states m)]
																else [0.0 | _<-list]
----------------------------------------------------																

----------------------------------------------------
-- This function implements the multiplication of a matrix and a vector, this is requiered for the X operator
next :: [[Float]] -> [Float] -> [Float]
next [] ys = []
next (xs:xss) ys = (mult xs ys) : next xss ys
				    where mult xs ys = sum (zipWith (*) xs ys)

-- Auxiliary function that returns a vector of 1 and 0's needed for computing the probability of the X operator
vecNext :: [Char] -> [Char] -> [Float]
vecNext [] _ = []
vecNext (x:xs) sat = if elem x sat 
						then 1.0:resto
						else 0.0:resto
					where resto = vecNext xs sat
----------------------------------------------------

----------------------------------------------------
-- Function that computes the set Sat^Yes
sYes :: Phi -> [Char]
sYes phi2 = [edo | edo<-(states mChain), elem edo (sat mChain phi2)]

-- Function that computes the set Sat^No
sNo :: Phi -> Phi -> [Char]
sNo phi1 phi2 = [edo | edo<-(states mChain), not(elem edo (sat mChain (O phi1 phi2)))]
----------------------------------------------------
-- Function that computes the set Sat^Yes for the unbounded until operator
sYesU :: Phi -> Phi-> [Char]
sYesU phi1 phi2 = (states mChain) \\ (prob0 (sat mChain phi1) (sat mChain phi2) [])

prob0 sat1 sat2 res = if sat2 == res 
                         then sat2
                         else prob0 sat1 l2 sat2 
  where
    l2 = union sat2 [(states mChain)!!i | i<-[0..length(states mChain)-1], 
                               elem ((states mChain)!!i) sat1,
                               or (map (>0) ((pTrans mChain)!!i))]  
----------------------------------------------------
----------------------------------------------------
-- Functions needed for the unbounded until, operator
-- implemented yet

diagMatrix :: [[Float]]
diagMatrix = [
				[1.0, 0.0,  0.0, 0.0],
				[0.0, 1.01, 0.0, 0.0],
				[0.0, 0.0,  1.0, 0.0],
				[0.0, 0.0,  0.0, 1.0]
				]

lu :: [[Float]]
lu = [
		[0.0, 0.0, 0.0,  0.0],	
		[0.0, 0.0, 0.01, 0.98],
		[0.0, 0.0, 0.0,  0.0],
		[0.0, 0.0, 0.0,  0.0]
		]

b :: [Float]
b = [0.0,0.0,0.0,1.0]

matrixMult :: [[Float]] -> [[Float]] -> [[Float]]
matrixMult xss yss = [mult xss ys | ys <- yss]
					where mult xss ys = [sum (zipWith (*) xs ys) | xs<-xss] 

inverseD::[[Float]] -> [[Float]]
inverseD xss = [invRow xs | xs <- xss]
				where invRow xs = [if x==0 then x else (1/x) | x<-xs]

jacobi::Integer->[Float]
jacobi 0 = [0.0, 0.5, 0.0, 0.5]
jacobi (n+1) = next	(diagMatrix)
					(zipWith (+) (next lu (jacobi n)) b)
